from __future__ import annotations

import base64
import datetime
import io
import os
import time
import uuid
from typing import Dict, Annotated, Any

import jwt
import uvicorn
from fastapi import FastAPI, Body, Request, HTTPException
from flask import jsonify
from qrcode.main import QRCode
from starlette.middleware.cors import CORSMiddleware
from starlette.responses import StreamingResponse, JSONResponse, HTMLResponse

from common.comm_utils import response_with_graph_agent_history, write2disk
from server.audio_server import text2speech
from server.image_server import create_image
from server.openai_server import create_server, clear_session_history
from server.video_server import check_image_for_video, generate_action_for_video, create_video_by_type, \
    generate_3d_role_for_video

server_instance = create_server()

app = FastAPI(docs_url=None, redoc_url=None, title="FunGPT", description="chatbot")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)
# app.mount("/static", StaticFiles(directory="static"), name="static")


# @app.get("/docs", include_in_schema=False)
# async def custom_swagger_ui_html():
#     return get_swagger_ui_html(
#         openapi_url=app.openapi_url,
#         title=app.title + " - Swagger UI",
#         oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
#         swagger_js_url="../static/swagger-ui-bundle.js",
#         swagger_css_url="/static/swagger-ui.css",
#     )


@app.post("/ai/query", summary="答问接口", tags=['doc'])
async def query(json_post: Annotated[Dict[str, Any], Body()], req: Request):
    """Handle a query request from the client.

        Returns:
            A JSON response containing the user ID, bot ID, message, response, conversation ID, and tags.
        """
    data = json_post
    user_id = data.get('user_id', '周佳乐')
    bot_id = data.get('bot_id', 'default_bot')
    message = data.get('message', '')
    conversation_id = data.get('conversation_id', 'default_conversation')
    tags = data.get('tags', [])
    input_json = {"question": message, "user_id": user_id}
    return StreamingResponse(response_with_graph_agent_history(user_id, message), media_type="text/event-stream")


@app.post("/ai/clear_chat_history", summary="清理聊天记录缓存接口", tags=['doc'])
async def query(json_post: Annotated[Dict[str, Any], Body()], req: Request):
    """Handle a query request from the client.

        Returns:
            A JSON response containing the user ID, bot ID, message, response, conversation ID, and tags.
        """
    data = json_post
    user_id = data.get('user_id', '周佳乐')
    session_id = data.get('session_id', 'LuckyZhou')
    message = data.get('message', '')
    clear_session_history(session_id)
    return JSONResponse(status_code=200, content="OK")


@app.post("/ai/handle_delete_entry", summary="答问接口", tags=['doc'])
async def handle_delete_entry(json_post: Annotated[Dict[str, Any], Body()], req: Request):
    """Handle a request to delete an entry from the knowledge base.

    Returns:
        A JSON response indicating the status of the deletion.
    """
    data = json_post
    entry_id = data.get('entry_id', '')

    if not entry_id:
        return jsonify({'status': 'error', 'message': 'Entry ID is required.'}), 400

    # Delete entry
    status = server_instance.knowledge_base.delete_entry(entry_id)

    if status:
        return jsonify({'status': 'success', 'message': 'Entry deleted successfully.'})
    else:
        return jsonify(
            {'status': 'error', 'message': 'Failed to delete entry. Please check the entry ID and try again.'}), 500


@app.post("/ai/create_images")
async def create_images(json_post: Annotated[Dict[str, Any], Body()], req: Request):
    try:
        description = json_post.get("description")
        negative_prompt = """
        nsfw,(worst quality, low quality:2),zombie,overexposure,text,bad anatomy,extra hands,((extra fingers)),too many fingers,fused fingers,((bad fingers)),bad arm,distorted arm,((extra arms)),fused arms,extra legs,missing leg,disembodied leg,extra nipples,detached arm,liquid hand,inverted hand,disembodied limb,small breasts,loli,oversized head,extra body,completely nude,extra navel,easynegative,(hair between eyes),sketch,duplicate,ugly,huge eyes,logo,worst face,(blurry:2),horror,geometry,bad_prompt,(bad hands),(missing fingers),multiple limbs,(interlocked fingers:1.2),Ugly Fingers,(extra digit and hands and fingers and legs and arms:1.4),(deformed fingers:1.2),(long fingers:1.2),bad-artist,(badhandv4:1.2),(By bad artist -neg),(bad-picture-chill-75v),sketches, normal quality,((monochrome)), ((grayscale)),skin spots,acnes,skin blemishes,(long hair:1.4),DeepNegative,(fat:1.2),facing away,looking away,tilted head,lowres,error,fewer digits,cropped,worstquality,jpegartifacts,username,bad feet,poorly drawn hands,poorly drawn face,mutation,jpeg artifacts,signature,watermark,extra limbs,malformed limbs,long neck,cross-eyed,mutated hands,polar,bad body,bad proportions,gross proportions,missing arms,extra foot,
        """
        # 是否开启prompt智能改写。开启后会使用大模型对输入prompt进行智能改写，仅对正向提示词有效。对于较短的输入prompt生成效果提升明显，但会增加3 - 4秒耗时。
        prompt_extend = json_post.get("prompt_extend",True)
        # 是否添加水印标识，水印位于图片右下角，文案为"AI生成"。
        watermark = json_post.get("watermark",False)
        sketch_image_url = json_post.get("sketch_image_url")
        # "https://help-static-aliyun-doc.aliyuncs.com/assets/img/zh-CN/6609471071/p743851.jpg"
        base_image_url = json_post.get("base_image_url")
        # "http://synthesis-source.oss-accelerate.aliyuncs.com/lingji/validation/mask2img/demo/source3.jpg"
        mask_image_url = json_post.get("base_image_url")
        # "http://synthesis-source.oss-accelerate.aliyuncs.com/lingji/validation/mask2img/demo/glasses.png"
        extra_input = None
        if base_image_url and sketch_image_url:
            extra_input = {
                "base_image_url": base_image_url,
                "mask_image_url": mask_image_url
            }
        result_data = create_image(description, negative_prompt=negative_prompt, sketch_image_url = sketch_image_url, extra_input = extra_input, prompt_extend = prompt_extend)
        return JSONResponse(status_code=200, content=result_data)


    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/ai/text_to_speech", summary="文字转语音接口", tags=['doc'])
async def text_to_speech(json_post: Annotated[Dict[str, Any], Body()]):
    """
    将文本转换为语音流
    """
    try:
        text = json_post.get("text", "")
        
        # 返回音频数据
        return StreamingResponse(
            # 调用 TTS 服务并收集所有音频数据
            io.BytesIO(text2speech(text)),
            media_type="audio/wav"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/ai/image2video_check", summary="检查图片是否适合生成视频", tags=['doc'])
async def image2video_check(json_post: Annotated[Dict[str, Any], Body()]):
    """
    检查提供的图片是否适合用于视频生成
    """
    try:
        image_url = json_post.get("image_url")
        # "animate-anyone-detect-gen2" 舞动人像
        type = json_post.get("type")
        if not type:
            raise HTTPException(status_code=400, detail="type参数错误")
        if not image_url:
            raise HTTPException(status_code=400, detail="缺少image_url参数")

        result = check_image_for_video(type=type, image_url=image_url, )
        return JSONResponse(status_code=200, content=result)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/ai/generate_video_action", summary="检查图片是否适合生成视频", tags=['doc'])
async def generate_video_action(json_post: Annotated[Dict[str, Any], Body()]):
    """
    提取视频动作
    """
    try:
        video_url = json_post.get("video_url")
        if not video_url:
            raise HTTPException(status_code=400, detail="缺少video_url参数")

        template_id = generate_action_for_video(video_url)
        return JSONResponse(status_code=200, content=template_id)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/ai/create_action_video_by_video", summary="舞动人像-根据图片和动作生成视频", tags=['doc'])
async def create_action_video_by_video(json_post: Annotated[Dict[str, Any], Body()]):
    """
    将图片人物实现视频动作
    """
    try:
        image_url = json_post.get("image_url")
        video_url = json_post.get("video_url")
        type = json_post.get("type")

        if type=="VIDEO":
            result = create_video_by_type(type=type, json_post=json_post)
            return JSONResponse(status_code=200, content=result)
        else:
            if type == "Motion":
                check_result = check_image_for_video(type=type, video_url=video_url)
            else:
                check_result = check_image_for_video(type=type, image_url=image_url, )

            if check_result["output"]["check_pass"]:
                if type == "ANIMATE" and video_url:
                    template_id = generate_action_for_video(video_url)
                    json_post["template_id"] = template_id
                elif type == "EMO":
                    face_bbox = check_result["output"]["face_bbox"]
                    ext_bbox = check_result["output"]["ext_bbox"]
                    if face_bbox and ext_bbox and ext_bbox:
                        json_post["face_bbox"] = face_bbox
                        json_post["ext_bbox"] = ext_bbox
                    else:
                        raise HTTPException(status_code=400, detail="EMO检测参数丢失")
                elif type == "Motion":
                    detected = check_result["output"]["detected"]
                    frame_index = check_result["output"]["frame_index"]
                    bbox = check_result["output"]["bbox"]
                    ply_url = generate_3d_role_for_video(image_url)
                    if detected and frame_index and bbox:
                        json_post["detected"] = detected
                        json_post["frame_index"] = frame_index
                        json_post["replacement_id"] = [ply_url]
                    else:
                        raise HTTPException(status_code=400, detail="Motion检测参数丢失")
                result = create_video_by_type(type=type, json_post=json_post)
                return JSONResponse(status_code=200, content=result)
            else:
                raise HTTPException(status_code=400, detail="图片不适合生成视频")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/images")
async def get_images(area: str = "山东省"):
    """获取指定区域的所有图片路径."""
    folder_path = f'lucky_web/public/photo/{area}'
    current_dir = os.path.dirname(os.path.abspath(__file__))
    full_folder_path = os.path.join(current_dir, folder_path)

    try:
        images = []
        for file in os.listdir(full_folder_path):
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')):
                # absolute_path = os.path.abspath(os.path.join(full_folder_path, file))
                # images.append(absolute_path)
                images.append(f"{area}/{file}")

        return JSONResponse(status_code=200, content=images)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/image_paths")
async def get_image_paths(area: str = "山东省"):
    """获取指定区域的所有图片路径."""
    folder_path = f'lucky_web/public/photo/{area}'
    current_dir = os.path.dirname(os.path.abspath(__file__))
    full_folder_path = os.path.join(current_dir, folder_path)

    try:
        images = []
        for file in os.listdir(full_folder_path):
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')):
                absolute_path = os.path.abspath(os.path.join(full_folder_path, file))
                images.append(absolute_path)

        return JSONResponse(status_code=200, content=images)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/visited_areas")
async def get_visited_areas():
    """获取所有访问过的区域."""
    folder_path = 'lucky_web/public/photo'
    current_dir = os.path.dirname(os.path.abspath(__file__))
    full_folder_path = os.path.join(current_dir, folder_path)

    try:
        areas = [name for name in os.listdir(full_folder_path) if
                 os.path.isdir(os.path.join(full_folder_path, name))]

        return JSONResponse(status_code=200, content=areas)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/common/download_video", summary="下载视频到本地", tags=['doc'])
async def download_video(json_post: Annotated[Dict[str, Any], Body()]):
    """
    下载视频到本地指定目录
    """
    try:
        video_url = json_post.get("video_url")
        if not video_url:
            raise HTTPException(status_code=400, detail="缺少video_url参数")
        
        # 可以从配置或环境变量中获取代理地址
        proxy = json_post.get("proxy") or os.getenv("HTTP_PROXY")
        save_path = "lucky_web/public/ai_material/ai_videos/"
        
        video_name = write2disk(video_url, save_path, proxy=proxy)

        return JSONResponse(status_code=200, content={
            "message": "视频下载成功",
            "file_path": save_path,
            "file_name": video_name
        })

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# JWT密钥
JWT_SECRET = "your_jwt_secret_key"

# 存储扫码场景值和状态
qr_scenes = {}

@app.get("/api/wx/qrcode")
async def get_wx_qrcode():
    """生成微信登录二维码"""
    scene = str(uuid.uuid4())
    qr_scenes[scene] = {
        "status": "WAITING",
        "created_at": time.time()
    }
    wx_app_id="wxa4e48ec87e373f87"
    domain_url="https://luckyzhou.online/"
    
    # 生成微信登录URL（需要替换为实际的微信开放平台参数）
    wx_url = f"https://open.weixin.qq.com/connect/oauth2/authorize?appid={wx_app_id}&redirect_uri={domain_url}&response_type=code&scope=snsapi_userinfo&state={scene}#wechat_redirect"
    
    # 生成二维码
    qr = QRCode(version=1, box_size=10, border=5)
    qr.add_data(wx_url)
    qr.make(fit=True)
    
    img = qr.make_image(fill_color="black", back_color="white")
    img_byte_arr = io.BytesIO()
    img.save(img_byte_arr, )
    img_byte_arr = img_byte_arr.getvalue()
    
    return JSONResponse(content={
        "qrcode": f"data:image/png;base64,{base64.b64encode(img_byte_arr).decode()}",
        "scene": scene
    })

@app.get("/api/wx/check_scan")
async def check_wx_scan(scene: str):
    """检查微信扫码状态"""
    if scene in qr_scenes:
        scene_data = qr_scenes[scene]
        if scene_data["status"] == "SCANNED":
            # 生成JWT token
            token = jwt.encode(
                {
                    "user_id": scene_data.get("user_id"),
                    "exp": datetime.datetime.now() + datetime.timedelta(days=1)
                },
                JWT_SECRET,
                algorithm="HS256"
            )
            
            # 清理场景值数据
            del qr_scenes[scene]
            
            return JSONResponse(content={
                "success": True,
                "token": token
            })
    
    return JSONResponse(content={"success": False})

# 微信回调接口
@app.get("/api/wx/callback")
async def wx_callback(code: str, state: str):
    """处理微信回调"""
    if state in qr_scenes:
        # # 通过code获取微信用户信息（需要实现）
        # user_info = await get_wx_user_info(code)
        user_info= {}
        qr_scenes[state].update({
            "status": "SCANNED",
            "user_id": user_info["openid"]
        })
        
        return HTMLResponse(content="登录成功，请返回网页")
    
    return HTMLResponse(content="登录失败，请重试")

@app.middleware("http")
async def verify_token(request: Request, call_next):
    """验证JWT token的中间件"""
    protected_paths = ["/api/images", "/api/image_paths", "/api/visited_areas"]
    
    if request.url.path in protected_paths:
        token = request.headers.get("Authorization")
        if not token:
            return JSONResponse(
                status_code=401,
                content={"detail": "Missing authentication token"}
            )
            
        try:
            # 验证token
            payload = jwt.decode(token.split(" ")[1], JWT_SECRET, algorithms=["HS256"])
            request.state.user = payload
        except jwt.ExpiredSignatureError:
            return JSONResponse(
                status_code=401,
                content={"detail": "Token has expired"}
            )
        except jwt.InvalidTokenError:
            return JSONResponse(
                status_code=401,
                content={"detail": "Invalid token"}
            )
    
    response = await call_next(request)
    return response

if __name__ == '__main__':
    uvicorn.run(app, host='0.0.0.0', port=8006, workers=1)
